##每次输出结果相同

import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
from sklearn.utils import shuffle

# 加载MNIST数据集，这里先把所有数据合并到一起
mnist = tf.keras.datasets.mnist
(X_all, y_all), _ = mnist.load_data()

# 数据归一化，明确转换数据类型为float32，确保后续操作不出问题
X_all = X_all.astype(np.float32) / 255.0
# 增加通道维度，使用更明确的reshape方式
X_all = X_all.reshape(X_all.shape[0], X_all.shape[1], X_all.shape[2], 1)

print("X_all shape:", X_all.shape)
y_all = y_all.reshape(-1, 1).astype(np.float32)  # 转换标签数据类型与图像数据类型一致，同时调整维度
print("y_all shape:", y_all.shape)

# 手动构建与X_all维度匹配的y_all_expanded
y_all_expanded = np.zeros((X_all.shape[0], X_all.shape[1], X_all.shape[2], 1), dtype=np.float32)
for i in range(X_all.shape[0]):
    y_all_expanded[i, :, :, 0] = y_all[i, 0]
print("y_all_expanded shape:", y_all_expanded.shape)

# 先将图像数据和标签数据分开处理，避免合并后出现难以处理的维度问题
train_images, test_images = np.split(X_all, [int(len(X_all) * 0.8)])
train_labels, test_labels = np.split(y_all, [int(len(y_all) * 0.8)])

print("train_images shape:", train_images.shape[0])
print("train_labels shape:", train_labels.shape[0])
print("test_images shape:", test_images.shape[0])
print("test_labels shape:", test_labels.shape[0])

model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_images, train_labels, epochs=5, validation_data=(test_images, test_labels))

# 模型评估
loss, accuracy = model.evaluate(test_images, test_labels)
print(f'Accuracy: {accuracy}')

# 预测
predictions = model.predict(test_images)

# 绘制前10个测试样本及其预测结果
for i in range(10):
    plt.subplot(2, 5, i + 1)
    plt.imshow(test_images[i].reshape(28, 28), cmap='gray')
    plt.title(f'Pred: {np.argmax(predictions[i])}\nTrue: {test_labels[i][0]}')
    plt.axis('off')
plt.show()